package poolx

import "sync"

type SizePool[T any] struct {
	sizedPool [64]*internalSizePool[T]
	selectF   func(size int) (index int)
	putF      func(data T) (index int)
}

func NewSizePool[T any](selectF func(size int) (index int), putF func(data T) (index int), newF func(size int) T) *SizePool[T] {
	ret := &SizePool[T]{selectF: selectF, putF: putF}
	for i := 0; i < 64; i++ {
		ret.sizedPool[i] = newInternalSizePool[T](newF, 1<<i)
	}
	return ret
}

func (s *SizePool[T]) Get(size int) T {
	index := s.selectF(size)
	index = index % 64
	pool := s.sizedPool[index]
	return pool.Get()
}

func (s *SizePool[T]) Put(data T) {
	index := s.putF(data)
	index = index % 64
	pool := s.sizedPool[index]
	pool.Put(data)
}

type internalSizePool[T any] struct {
	pool sync.Pool
}

func newInternalSizePool[T any](newF func(size int) T, size int) *internalSizePool[T] {
	return &internalSizePool[T]{
		pool: sync.Pool{
			New: func() any {
				return newF(size)
			},
		},
	}
}

func (s *internalSizePool[T]) Get() T {
	return s.pool.Get().(T)
}

func (s *internalSizePool[T]) Put(v T) {
	s.pool.Put(v)
}
